import argparse
import copy
import os
from obs import ObsClient

sk = "1lNcDkVMBV7tVQq3aBWXirTWClbpiWJ1r08zbkSD"
ak = "HPUABVUGOTP2OPODPEKP"
server = "https://obs.cn-southwest-2.myhuaweicloud.com"
bucket_name = "bucket-omni-infer-gy"



class ObsTool:
    def __init__(self, pipeline_name, pipeline_number):
        object_key_prefix = "CI_Pipeline/"
        self.object_key_prefix = f"{object_key_prefix}{pipeline_name}/{pipeline_number}/"


    def upload(self, file_path: str, object_key: str):
        upload_path = object_key
        if upload_path.startswith("/"):
            upload_path = upload_path[1:]

        obs_client = ObsClient(access_key_id=ak, secret_access_key=sk, server=server)
        for i in range(3):
            try:
                object_key = self.object_key_prefix + upload_path
                resp = obs_client.putFile(bucket_name, object_key, file_path, progressCallback=callback)
                print(f"upload resp: {resp}")
                print(f"object_key: {object_key}")
                if resp.status < 300:
                    print(f"file uploaded success: {object_key}")
                    break
                else:
                    print(f"failed to upload file. Status code: {resp.errorCode}, {resp.requestId}")
                    print(f"retry upload: {i + 1}")
            except Exception as exception:
                print(f"except to upload file: {exception}")
                print(f"retry upload: {i + 1}")
        obs_client.close()

    def upload_folder(self, file_path: str, object_key: str):
        obs_path = object_key
        if obs_path.endswith("/"):
            obs_path = obs_path[:-1]

        folder_path = file_path
        if folder_path.endswith("/") or folder_path.endswith("\\"):
            folder_path = folder_path[:-1]

        for root, dirs, files in os.walk(folder_path):
            for file in files:
                local_file_path = os.path.join(root, file)
                local_file_path = local_file_path.replace("\\", "/")
                obs_file_path = obs_path + local_file_path.replace(folder_path, '')
                self.upload(local_file_path, obs_file_path)

    def get_download_url(self, object_key: str):
        download_path = object_key
        if download_path.startswith("/"):
            download_path = download_path[1:]

        obs_client = ObsClient(access_key_id=ak, secret_access_key=sk, server=server)
        try:
            resp = obs_client.createSignedUrl(method='GET', bucketName=bucket_name, objectKey=self.object_key_prefix + download_path, expires=3600*18)
            print(resp.signedUrl)
            return resp.signedUrl
        except Exception as exception:
            print(f"except to get download url: {exception}")


def callback(transferred_amount, total_amount, total_seconds):
    # 获取上传平均速率(KB/S)
    print(transferred_amount * 1.0 / total_seconds / 1024)
    # 获取上传进度百分比
    print(transferred_amount * 100.0 / total_amount)

    #TODO 没返回值？？？


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Parameters of obs_util',
                                    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--type", type=str,
                        help="操作类型：upload或download", default="upload")
    parser.add_argument("--object_key", type=str,
                        help="obs文件路径", default="/test")
    parser.add_argument("--local_path", type=str, required=False,
                        help="本地文件（文件夹）路径", default=r"D:\pangu\server.log")
    parser.add_argument("--pipeline_name", type=str,
                        help="流水线名称", default="test")
    parser.add_argument("--pipeline_number", type=str,
                        help="流水线执行编号", default="2")

    result, _ = parser.parse_known_args()
    args = copy.deepcopy(result)

    _type = args.type
    _object_key = args.object_key
    _local_path = args.local_path
    _pipeline_name = args.pipeline_name
    _pipeline_number = args.pipeline_number
    # object_key_prefix = object_key_prefix + _pipeline_name + "/" + _pipeline_number + "/"
    obs_tool = ObsTool(_pipeline_name, _pipeline_number)

    if _type != 'upload' and _type != 'download':
        print("type param error")
        exit()

    if _type == 'upload':
        if _local_path is None:
            print("local_path is required")
            exit()
        if os.path.isdir(_local_path):
            obs_tool.upload_folder(_local_path, _object_key)
        else:
            obs_tool.upload(_local_path, _object_key)
    else:
        obs_tool.get_download_url(_object_key)



